use core::intrinsics::assume;
use core::{marker::PhantomData, ptr::NonNull};

use crate::addrspace::{PfnAddr, PhysAddr};

use super::{Memblock, MemblockFlags, MemblockRegion, MemblockType};

impl Memblock {
    /// 返回遍历预留域的迭代器
    pub fn iter_reserve(&self) -> IterRes {
        IterRes::new(self)
    }
    /// 返回遍历空闲内存域的迭代器 (避开预留内存)
    pub fn iter_free(&self, flags: MemblockFlags) -> IterFree {
        IterFree::new(self, flags)
    }
    /// 返回反向遍历空闲内存域的迭代器 (避开预留内存)
    pub fn iter_free_rev(&self, flags: MemblockFlags) -> IterFreeRev {
        IterFreeRev::new(self, flags)
    }
    /// 返回遍历可用内存 pfn 的迭代器
    pub fn iter_mem_pfn(&self) -> IterMemPfn {
        IterMemPfn::new(self)
    }
    /// 返回遍历可用内存的迭代器
    pub fn iter_each_mem(&self) -> IterMem {
        IterMem::new(self)
    }
    /// 返回遍历可用内存的可变迭代器 (仅该 crate 内部可用)
    pub(super) fn iter_mut_each_mem(&mut self) -> IterMemMut<'_> {
        IterMemMut::new(self)
    }
}

/// 遍历预留域的迭代器
///
/// 迭代器返回元组 (start, end)
/// * start: 当前域的起始地址
/// * end: 当前域的结束地址
pub struct IterRes<'a> {
    inner: &'a Memblock,
    idx: usize,
}

impl<'a> IterRes<'a> {
    fn new(mb: &'a Memblock) -> Self {
        Self { inner: mb, idx: 0 }
    }
}

impl<'a> Iterator for IterRes<'a> {
    type Item = (PhysAddr, usize);

    fn next(&mut self) -> Option<Self::Item> {
        if self.idx < self.inner.reserved.cnt {
            let r = self.inner.reserved.regions[self.idx];
            self.idx += 1;
            Some((r.base, r.size))
        } else {
            None
        }
    }
}

/// 遍历空闲内存域的迭代器 (避开预留内存)
///
/// 迭代器返回元组 (start, end)
/// * start: 当前域的起始地址
/// * end: 当前域的结束地址
pub struct IterFree<'a> {
    idx_a: usize,
    idx_b: usize,
    inner: &'a Memblock,
    flags: MemblockFlags,
}

impl<'a> IterFree<'a> {
    fn new(mb: &'a Memblock, flags: MemblockFlags) -> Self {
        Self {
            idx_a: 0,
            idx_b: 0,
            inner: mb,
            flags,
        }
    }
}

impl<'a> Iterator for IterFree<'a> {
    type Item = (PhysAddr, PhysAddr);

    fn next(&mut self) -> Option<Self::Item> {
        let tulpe = __next_mem_range(
            self.idx_a,
            self.idx_b,
            self.flags,
            &self.inner.memory,
            &self.inner.reserved,
        );

        if tulpe.0 == usize::MAX && tulpe.1 == usize::MAX {
            None
        } else {
            self.idx_a = tulpe.0;
            self.idx_b = tulpe.1;
            Some((tulpe.2, tulpe.3))
        }
    }
}

fn __next_mem_range(
    idx_a: usize,
    idx_b: usize,
    flags: MemblockFlags,
    type_a: &MemblockType,
    type_b: &MemblockType,
) -> (usize, usize, PhysAddr, PhysAddr) {
    for mut i in idx_a..type_a.cnt {
        let m = type_a.regions[i];

        let m_start = m.base;
        let m_end = m.base + m.size;
        if flags != m.flags {
            continue;
        }

        for mut j in idx_b..type_b.cnt + 1 {
            let r_start = if j != 0 {
                type_b.regions[j - 1].base + type_b.regions[j - 1].size
            } else {
                PhysAddr::from(0)
            };
            let r_end = if j < type_b.cnt {
                type_b.regions[j].base
            } else {
                PhysAddr::from(usize::MAX)
            };

            if r_start >= m_end {
                break;
            }
            if m_start < r_end {
                if m_end <= r_end {
                    i += 1;
                } else {
                    j += 1;
                }
                return (i, j, m_start.max(r_start), m_end.min(r_end));
            }
        }
    }
    (usize::MAX, usize::MAX, PhysAddr::from(0), PhysAddr::from(0))
}

/// 反向遍历空闲内存域的迭代器 (避开预留内存)
///
/// 迭代器返回元组 (start, end)
/// * start: 当前域的起始地址
/// * end: 当前域的结束地址
pub struct IterFreeRev<'a> {
    idx_a: usize,
    idx_b: usize,
    inner: &'a Memblock,
    flags: MemblockFlags,
}

impl<'a> IterFreeRev<'a> {
    fn new(mb: &'a Memblock, flags: MemblockFlags) -> Self {
        Self {
            idx_a: usize::MAX,
            idx_b: usize::MAX,
            inner: mb,
            flags,
        }
    }
}

impl<'a> Iterator for IterFreeRev<'a> {
    type Item = (PhysAddr, PhysAddr);

    fn next(&mut self) -> Option<Self::Item> {
        let tulpe = __next_mem_range_rev(
            self.idx_a,
            self.idx_b,
            self.flags,
            &self.inner.memory,
            &self.inner.reserved,
        );

        if tulpe.0 == usize::MAX && tulpe.1 == usize::MAX {
            None
        } else {
            self.idx_a = tulpe.0;
            self.idx_b = tulpe.1;
            Some((tulpe.2, tulpe.3))
        }
    }
}

fn __next_mem_range_rev(
    mut idx_a: usize,
    mut idx_b: usize,
    flags: MemblockFlags,
    type_a: &MemblockType,
    type_b: &MemblockType,
) -> (usize, usize, PhysAddr, PhysAddr) {
    if idx_a == usize::MAX && idx_b == usize::MAX {
        idx_a = type_a.cnt - 1;
        idx_b = type_b.cnt;
    }

    for mut i in (0..idx_a + 1).rev() {
        let m = type_a.regions[i];
        let m_start = m.base;
        let m_end = m.base + m.size;
        if flags != m.flags {
            continue;
        }

        for mut j in (0..idx_b + 1).rev() {
            let r_start = if j != 0 {
                type_b.regions[j - 1].base + type_b.regions[j - 1].size
            } else {
                PhysAddr::from(0)
            };
            let r_end = if j < type_b.cnt {
                type_b.regions[j].base
            } else {
                PhysAddr::from(usize::MAX)
            };

            if r_end <= m_start {
                break;
            }

            if m_end > r_start {
                if m_start >= r_start {
                    i += 1;
                } else {
                    j += 1;
                }
                return (i, j, m_start.max(r_start), m_end.min(r_end));
            }
        }
    }
    (usize::MAX, usize::MAX, PhysAddr::from(0), PhysAddr::from(0))
}

/// 遍历可用内存 pfn 的迭代器
///
/// 迭代器返回元组 (start_pfn, end_pfn)
/// * start_pfn: 当前域的 pfn 起始
/// * end_pfn: 当前域的 pfn 结束
pub struct IterMemPfn<'a> {
    inner: &'a Memblock,
    idx: i32,
}

impl<'a> IterMemPfn<'a> {
    fn new(mb: &'a Memblock) -> Self {
        Self { inner: mb, idx: -1 }
    }
}

impl<'a> Iterator for IterMemPfn<'a> {
    type Item = (PfnAddr, PfnAddr);

    fn next(&mut self) -> Option<Self::Item> {
        let tulpe = __next_mem_pfn_range(self.idx, &self.inner.memory);

        if tulpe.0 == -1 {
            None
        } else {
            self.idx = tulpe.0;
            Some((tulpe.1, tulpe.2))
        }
    }
}

fn __next_mem_pfn_range(mut idx: i32, ty: &MemblockType) -> (i32, PfnAddr, PfnAddr) {
    let mut start = PfnAddr::from(0);
    let mut end = PfnAddr::from(0);
    idx += 1;

    while idx < ty.cnt as i32 {
        let r = ty.regions[idx as usize];

        start = r.base.phys_to_pfn_up();
        end = (r.base + r.size).phys_to_pfn_down();
        if start >= end {
            idx += 1;
            continue;
        }
        break;
    }
    if idx >= ty.cnt as i32 {
        return (-1, PfnAddr::from(0), PfnAddr::from(0));
    }
    (idx, start, end)
}

/// 遍历可用内存的迭代器
///
/// 迭代器返回元组 (start, size, flags)
/// * start: 当前域物理内存的起始地址
/// * size: 当前域的大小
/// * flags: 当前域的标识
pub struct IterMem<'a> {
    inner: &'a Memblock,
    idx: usize,
}

impl<'a> IterMem<'a> {
    fn new(mb: &'a Memblock) -> Self {
        Self { inner: mb, idx: 0 }
    }
}

impl<'a> Iterator for IterMem<'a> {
    type Item = (PhysAddr, usize, MemblockFlags);

    fn next(&mut self) -> Option<Self::Item> {
        if self.idx < self.inner.memory.cnt {
            let i = self.idx;

            self.idx += 1;
            let rgn = self.inner.memory.regions[i];
            Some((rgn.base, rgn.size, rgn.flags))
        } else {
            None
        }
    }
}

/// 遍历可用内存的可变迭代器
///
/// 迭代器返回对每一个 MemblockRegion 的可变引用
pub(super) struct IterMemMut<'a> {
    ptr: NonNull<MemblockRegion>,
    end: *mut MemblockRegion,
    _marker: PhantomData<&'a mut MemblockRegion>,
}

impl<'a> IterMemMut<'a> {
    fn new(mb: &'a mut Memblock) -> Self {
        let ptr = mb.memory.regions.as_mut_ptr();

        unsafe {
            assume(!ptr.is_null());

            let end = ptr.add(mb.memory.cnt);

            Self {
                ptr: NonNull::new_unchecked(ptr),
                end,
                _marker: PhantomData,
            }
        }
    }
}

impl<'a> Iterator for IterMemMut<'a> {
    type Item = &'a mut MemblockRegion;

    fn next(&mut self) -> Option<Self::Item> {
        unsafe {
            assume(!self.end.is_null());
            if self.ptr.as_ptr() as *const MemblockRegion == self.end {
                None
            } else {
                let old = self.ptr.as_ptr();
                self.ptr = NonNull::new_unchecked(self.ptr.as_ptr().add(1));
                Some(&mut (*old))
            }
        }
    }
}
